import numpy as np
cimport numpy as np
import cPickle as pickle
cimport cython

from libc.math cimport sqrt

ctypedef np.float64_t flt_t
ctypedef np.complex128_t cpx_t
#ctypedef np.int_t nt_t

flt = np.float64
cpx = np.complex128
nt = np.int


#@cython.profile(False)
cdef flt_t heaviside(flt_t x):
    if x == 0.:
        return 0.5
    elif x < 0.:
        return 0.
    else: 
        return 1.

#@cython.profile(False)
cdef flt_t interp(flt_t t, np.ndarray[flt_t, ndim=1] tvec,
                     np.ndarray[flt_t, ndim=1] yvec):
    #find max index where tvec<t
    cdef int tvec_idx, tvec_len = len(tvec), max_idx=len(tvec) - 1
    cdef flt_t result = yvec[max_idx]

    for tvec_idx in range(tvec_len):
        if tvec[tvec_idx] > t:
            max_idx = tvec_idx - 1
            result = yvec[max_idx] + (yvec[max_idx + 1]- yvec[max_idx]) * \
                (t - tvec[max_idx]) / (tvec[max_idx + 1] - tvec[max_idx])
            break
    return result

#@cython.profile(False)
cdef np.ndarray fieldeqns(np.ndarray[cpx_t, ndim=1] y, flt_t t,
                np.ndarray[flt_t, ndim=1] epsilont, 
                np.ndarray[flt_t, ndim=1] epsilon, flt_t chi,
                flt_t kappa1, flt_t kappa2, flt_t Delta_a,
                flt_t Delta_b, int a, int b, int c):
    
    cdef np.ndarray[cpx_t, ndim=1] dy
    dy = np.zeros([2], dtype=cpx)
    
    chi = -chi * ((-1.) ** a + (-1.) ** b + (-1.) ** c)
    eps = interp(t, epsilont, epsilon)
    
    # Interpolate the data set (gt,g) at time t
    dy[0] = -sqrt(kappa1) * eps - 1j * (Delta_a + chi) * y[0] - \
             0.5 * kappa1 * y[0] - 0.5 * sqrt(kappa1 * kappa2) * y[1]
    dy[1] = -sqrt(kappa2) * eps - 1j * (Delta_b + chi) * y[1] - \
             0.5 * kappa2 * y[1] - 0.5 * sqrt(kappa1 * kappa2)*y[0]
    return dy

#@cython.profile(False)
cdef inline flt_t chi_f(flt_t chi, int a, int b, int c):
    return -chi*((-1.)**a + (-1.)**b + (-1.)**c)

cdef inline int imod(int a, int b):
    #gives a mod b for a, b positive integers
    cdef int c = a
    while a > b:
        c = c - b
    return c

@cython.boundscheck(False)
def get_trajectory(dict input_dict):
    """
    Produces a numerical solution to the stochastic system, given a
    series of parameters from a dictionary pickled in an input file.
    Parameters include:
     + init: the initial density matrix
     + T: total simulation time
     + dt: time stepping
     + kappa1: cavity decay rate for mode 1
     + kappa2: cavity decay rate for mode 2
     + epsilon: an array containing the measurement pulse
     + chi: cavity-qubit coupling constant
     + tmeas_on: time at which measurement is initiated
     + tmeas_off: time at which measurement stops.
     + Delta_a: Field equation parameter
     + Delta_b: Field equation parameter
     + D1: time-independent deterministic part of liouvillian
    """
    
    #with open(input_file, 'r') as file:
    #    input_dict = pickle.load(file)
    
    #cdef np.ndarray[cpx_t, ndim=1]  
    #Hardcode initial state to avoid side effect, 'init' no longer nec.
    cdef np.ndarray[cpx_t, ndim=1] rho = 0.125 * np.ones(64, dtype=cpx)
    cdef np.ndarray[flt_t, ndim=1] tvec = input_dict['tvec'] 
    cdef flt_t kappa1 = input_dict['kappa1']
    cdef flt_t kappa2 = input_dict['kappa2']
    cdef flt_t chi = input_dict['chi']
    cdef flt_t eta = input_dict['eta']
    cdef flt_t tmeas_on = input_dict['tmeas_on']
    cdef flt_t tmeas_off = input_dict['tmeas_off']
    cdef np.ndarray[cpx_t, ndim=2] D1 = input_dict['D1']
    cdef np.ndarray[cpx_t, ndim=1] ideal_p = input_dict['ideal_p']
    cdef np.ndarray[cpx_t, ndim=1] ideal_m = input_dict['ideal_m']
    cdef np.ndarray[flt_t, ndim=1] epsilon = input_dict['epsilon']
    cdef flt_t Delta_a = input_dict['Delta_a']
    cdef flt_t Delta_b = input_dict['Delta_b']
    cdef int N = input_dict['N']
    cdef int dim = input_dict['dim']

    cdef dict output_dict
    
    #Static Typing for all loop variables:
    cdef int idx, indx, a, b, c, d, e, f, lin_idx, m, n, m_a, m_b, n_a
    cdef int n_b, l, r
    
    cdef flt_t dt, t, wiener_increment, kernel_mean=0.
    cdef np.ndarray[flt_t, ndim=1] filter_kernel, dW
    cdef np.ndarray[cpx_t, ndim=1] sigma000, sigma011
    cdef np.ndarray[cpx_t, ndim=2] sol
    cdef np.ndarray[cpx_t, ndim=4] alpha, beta
    cdef np.ndarray[cpx_t, ndim=1] temp_slice #increases efficiency of ODE solution
    cdef np.ndarray[cpx_t, ndim=1] D1_t #recent change
    cdef np.ndarray[cpx_t, ndim=1] Am
    cdef np.ndarray[cpx_t, ndim=1] D2 
    cdef np.ndarray[cpx_t, ndim=1] d1, d2, zt, zp, zm, d1t
    cdef cpx_t j, s, s_filtered, purity
    cdef int am_flag = 1

    dt = tvec[1]-tvec[0]
    
    # Solve for cavity state, filter kernel
    alpha = np.zeros(shape=(N, 2, 2, 2), dtype=cpx)
    beta = np.zeros(shape=(N, 2, 2, 2), dtype=cpx)
    
    #Explicit RK4:
    for indx in range(8):
        a, b, c = int(indx > 3), int(indx % 4 > 1), indx % 2 #bits
        
        sol = np.zeros((N, 2), dtype=cpx)
        t = 0.
        for idx in range(1,N):
            temp_slice = np.array([sol[idx-1, 0], sol[idx-1, 1]])
            k1 = fieldeqns(temp_slice, t,
                            tvec, epsilon, chi, kappa1, kappa2,
                            Delta_a, Delta_b, a, b, c)
            k2 = fieldeqns(temp_slice + 0.5 * dt * k1, t + 0.5 * dt,
                            tvec, epsilon, chi, kappa1, kappa2, 
                            Delta_a, Delta_b, a, b, c)
            k3 = fieldeqns(temp_slice + 0.5 * dt * k2, t + 0.5 * dt,
                            tvec, epsilon, chi, kappa1, kappa2, 
                            Delta_a, Delta_b, a, b, c)
            k4 = fieldeqns(temp_slice + dt * k3, t + dt,
                            tvec, epsilon, chi, kappa1, kappa2, 
                            Delta_a, Delta_b, a, b, c)
            sol[idx, 0] = temp_slice[0] + dt / 6. * (k1[0] + 2. * k2[0] + 2. * k3[0] + k4[0])
            sol[idx, 1] = temp_slice[1] + dt / 6. * (k1[1] + 2. * k2[1] + 2. * k3[1] + k4[1])
            t += dt
        for lin_idx in range(N):
            alpha[lin_idx, a, b, c] = sol[lin_idx, 0]
            beta[lin_idx, a, b, c] = sol[lin_idx, 1]


    sigma000 = np.zeros(N, dtype=cpx)
    sigma011 = np.zeros(N, dtype=cpx)
    for lin_idx in range(N):
        sigma000[lin_idx] = sqrt(kappa1) * alpha[lin_idx, 0, 0, 0] + \
                             sqrt(kappa2) * beta[lin_idx, 0, 0, 0]
        sigma011[lin_idx] = sqrt(kappa1) * alpha[lin_idx, 0, 1, 1] + \
                             sqrt(kappa2) * beta[lin_idx, 0, 1, 1]
    filter_kernel = np.zeros(N, dtype=flt)
    for lin_idx in range(N):
        filter_kernel[lin_idx] = -(sigma000[lin_idx].imag +
                                             3. * sigma011[lin_idx].imag)
        kernel_mean = kernel_mean + filter_kernel[lin_idx]/N
    for lin_idx in range(N):
        filter_kernel[lin_idx] = filter_kernel[lin_idx] / kernel_mean    

    #Values for main loop to fill/increment
    s = 0.
    s_filtered = 0.
    
    dW = np.random.randn(N)
    Am = np.zeros(shape=(dim,), dtype=cpx)
    D1_t = np.zeros(shape=(dim ** 2,), dtype=cpx)
    D2 = np.zeros((dim**2,),dtype=cpx)
    d1 = np.zeros((dim**2,),dtype=cpx)
    d2 = np.zeros((dim**2,),dtype=cpx)
    d2p = np.zeros((dim**2,),dtype=cpx)
    d2m = np.zeros((dim**2,),dtype=cpx)
    zt = np.zeros((dim**2,),dtype=cpx)
    zp = np.zeros((dim**2,),dtype=cpx)
    zm = np.zeros((dim**2,),dtype=cpx)
    d1t = np.zeros((dim**2,),dtype=cpx)
    # Time stepping using Kloeden-Platen method with fixed stepsize dt
    
    #Array Observables
    #Actual Parity, time dependent measurement operator, photocurrent
    cdef np.ndarray[cpx_t, ndim=1] tr_rho_ZZZ, tr_rho_pi_sigma
    cdef np.ndarray[cpx_t, ndim=1] photocurrent
    tr_rho_ZZZ = np.zeros((N,),dtype=cpx)
    tr_rho_pi_sigma = np.zeros((N,),dtype=cpx)
    photocurrent = np.zeros((N,),dtype=flt)
    
    for idx in range(N):
        wiener_increment = sqrt(dt) * dW[idx]
        
        # the time dependent measurement operator
        if tvec[idx] > tmeas_on and tvec[idx] < tmeas_off:
            for indx in range(dim):
                a, b, c = int(indx > 3), int(indx % 4 > 1), indx % 2
                Am[indx] = 1j * (sqrt(kappa1 * eta) * alpha[idx, a, b, c]
                + sqrt(kappa2 * eta) * beta[idx, a, b, c])
        elif am_flag == 1:
            for indx in range(dim):
                Am[indx] = 0.0
            am_flag = 0 #ensure Am is only reset to 0 once.

        for m in range(dim**2):
            #To take kronecker product, get little indices from big indices:
            m_b = m % dim
            m_a = (m - m_b)/dim
            D2[m] = Am[m_a] + Am[m_b].conjugate()

        average = 0.
        for indx in range(dim):
            average = average + \
            2. * Am[indx].real * rho[indx * (dim + 1)]

        # the current
        j = average + wiener_increment / dt
        photocurrent[idx] = j
        # the integrated current
        s = s + j * dt
        
        # the integrated filtered current
        s_filtered = s_filtered + filter_kernel[idx] * \
                        j * dt

        #Parity
        tr_rho_ZZZ[idx] = rho[0] - rho[9] - rho[18] + rho[27] - \
                            rho[36] + rho[45] + rho[54] - rho[63]

        #Measured Qubit operator
        for indx in range(dim):
            a, b, c = int(indx > 3), int(indx % 4 > 1), indx % 2
            tr_rho_pi_sigma[idx] = tr_rho_pi_sigma[idx] +\
                (sqrt(kappa1) * alpha[idx, a, b, c] + 
                    sqrt(kappa2) * beta[idx, a, b, c]) * \
                rho[dim * indx + indx] 

        # The time dependent part of the deterministic super operator
        
        for indx in range(dim**2):
            a, b, c, d, e, f = int(indx>31), int(indx%32>15), int(indx%16>7), \
                               int(indx%8>3), int(indx%4>1), indx%2
            l, r = (indx - indx%8)/8, indx%8
            
            D1_t[dim*l + r] = ( (chi_f(chi, d, e, f) - chi_f(chi, a, b, c)) * (alpha[idx, a, b, c].conjugate() * alpha[idx, d, e, f]).imag + \
                                              (chi_f(chi, d, e, f) - chi_f(chi, a, b, c)) * (beta[idx, a, b, c].conjugate() * beta[idx, d, e, f]).imag)
        for indx in range(dim):
            a, b, c = int(indx > 3), int(indx % 4 > 1), indx % 2
            d, e, f = 1-a, 1-b, 1-c
            
            D1_t[7 * indx + 7] = D1_t[7 * indx + 7] + sqrt(kappa1*kappa2) * \
                                            ( ( beta[idx, a, b, c] * (alpha[idx, d, e, f].conjugate() )).real + \
                                              (alpha[idx, a, b, c] * ( beta[idx, d, e, f].conjugate() )).real - \
                                              ( beta[idx, a, b, c] * (alpha[idx, a, b, c].conjugate() )).real - \
                                              ( beta[idx, d, e, f] * (alpha[idx, d, e, f].conjugate() )).real )
        
        # Order 2 weak Runge-Kutta scheme (Platen)
        # Temporarily modify D1:
        for indx in range(dim**2):
            D1[indx, indx] = D1[indx, indx] + D1_t[indx]
        d1 = np.dot(D1, rho)
        # D1(psi)
        for indx in range(dim**2):
            d2[indx] = (D2[indx] - average) * rho[indx]
        # D2(psi)
        for indx in range(dim**2):
            zt[indx] = rho[indx] + d1[indx] * dt + d2[indx] * wiener_increment
        # psi_tilde
        for indx in range(dim**2):
            zp[indx] = rho[indx] + d1[indx] * dt + d2[indx] * sqrt(dt)
        # psi_plus
        for indx in range(dim**2):
            zm[indx] = rho[indx] + d1[indx] * dt - d2[indx] * sqrt(dt)
        # psi_minus
        d1t = np.dot(D1, zt)
        #Undo temporary D1 mod
        for indx in range(dim**2):
            D1[indx, indx] = D1[indx, indx] - D1_t[indx]
        # D1(psi)
        average_p = 0.
        average_m = 0.
        for indx in range(dim):
            average_p = average_p + \
            2.* Am[indx].real * zp[indx * (dim + 1)]
            average_m = average_m + \
            2.* Am[indx].real * zm[indx * (dim + 1)]
        for indx in range(dim**2):
            d2p[indx] = (D2[indx] - average_p) * zp[indx]
            d2m[indx] = (D2[indx] - average_m) * zm[indx]
        
        rho += 0.5 * ( d1t + d1 ) * dt + 0.25 * ( d2p + d2m + 2.*d2 ) * wiener_increment\
        + 0.25 * ( d2p - d2m ) * ( wiener_increment ** 2 - dt ) / sqrt(dt)
    
    purity = 0.
    for indx in range(dim):
        for idx in range(dim):
            #purity = sum(rho_jk * rho_kj) = Tr(rho^2).
            purity = purity + \
                        rho[dim * indx + idx] * rho[dim * idx + indx]

    output_dict = {'s':s, 's_filtered':s_filtered,
                    'overlap_p':np.dot(rho.conj().T,ideal_p).real, 
                    'overlap_m':np.dot(rho.conj().T,ideal_m).real,
                    'purity':purity, 'tr_rho_ZZZ': tr_rho_ZZZ, 
                    'tr_rho_pi_sigma': tr_rho_pi_sigma}

    return output_dict
